package com.hanserwei.han_picture.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.hanserwei.han_picture.domain.entity.po.User;
import com.hanserwei.han_picture.manager.RedisManager;
import com.hanserwei.han_picture.websocket.message.AiProcessMessage;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.lang.NonNull;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.net.URI;
import java.util.concurrent.ConcurrentHashMap;

/**
 * AI处理WebSocket处理器
 *
 * @author hanserwei
 */
@Component
@Slf4j
@RequiredArgsConstructor
public class AiProcessWebSocketHandler implements WebSocketHandler {

    private final RedisManager redisManager;
    private final ObjectMapper objectMapper;
    
    /**
     * 存储用户ID与WebSocket会话的映射
     * key: userId, value: WebSocketSession
     */
    private final ConcurrentHashMap<Long, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(@NonNull WebSocketSession session) throws Exception {
        log.info("WebSocket连接建立: {}", session.getId());

        // 从查询参数中获取用户ID
        Long userId = getUserIdFromSession(session);
        if (userId != null) {
            // 验证用户身份
            if (validateUser(userId)) {
                userSessions.put(userId, session);
                log.info("用户 {} 的WebSocket连接已建立", userId);

                // 发送连接成功消息
                sendMessage(session, AiProcessMessage.createConnectionSuccessMessage(userId));
            } else {
                log.warn("用户身份验证失败，关闭连接: userId={}", userId);
                session.close(CloseStatus.NOT_ACCEPTABLE.withReason("用户身份验证失败"));
            }
        } else {
            log.warn("无法获取用户ID，关闭连接");
            session.close(CloseStatus.NOT_ACCEPTABLE.withReason("缺少用户ID"));
        }
    }

    @Override
    public void handleMessage(@NonNull WebSocketSession session, @NonNull WebSocketMessage<?> message) throws Exception {
        log.debug("收到WebSocket消息: {}", message.getPayload());

        // 处理心跳消息
        if ("ping".equals(message.getPayload().toString())) {
            sendMessage(session, AiProcessMessage.createHeartbeatMessage());
        }
    }

    @Override
    public void handleTransportError(@NonNull WebSocketSession session, @NonNull Throwable exception) throws Exception {
        log.error("WebSocket传输错误: {}", session.getId(), exception);
        removeSession(session);
    }

    @Override
    public void afterConnectionClosed(@NonNull WebSocketSession session, @NonNull CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接关闭: {}, 状态: {}", session.getId(), closeStatus);
        removeSession(session);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 向指定用户发送AI处理消息
     */
    public void sendAiProcessMessage(Long userId, AiProcessMessage message) {
        WebSocketSession session = userSessions.get(userId);
        if (session != null && session.isOpen()) {
            sendMessage(session, message);
        } else {
            log.warn("用户 {} 的WebSocket连接不存在或已关闭，无法发送消息", userId);
        }
    }

    /**
     * 发送消息到WebSocket会话
     */
    private void sendMessage(WebSocketSession session, AiProcessMessage message) {
        try {
            String messageJson = objectMapper.writeValueAsString(message);
            session.sendMessage(new TextMessage(messageJson));
            log.debug("发送WebSocket消息成功: {}", messageJson);
        } catch (Exception e) {
            log.error("发送WebSocket消息失败: {}", session.getId(), e);
            removeSessionBySession(session);
        }
    }

    /**
     * 从会话中获取用户ID
     */
    private Long getUserIdFromSession(WebSocketSession session) {
        try {
            URI uri = session.getUri();
            if (uri != null) {
                String query = uri.getQuery();
                if (query != null) {
                    String[] params = query.split("&");
                    for (String param : params) {
                        String[] keyValue = param.split("=");
                        if (keyValue.length == 2 && "userId".equals(keyValue[0])) {
                            return Long.parseLong(keyValue[1]);
                        }
                    }
                }
            }
        } catch (Exception e) {
            log.error("解析用户ID失败", e);
        }
        return null;
    }

    /**
     * 验证用户身份
     */
    private boolean validateUser(Long userId) {
        try {
            // 从Redis中获取用户会话验证身份
            Object userSession = redisManager.getUserSession(userId);
            return userSession instanceof User;
        } catch (Exception e) {
            log.error("验证用户身份失败: userId={}", userId, e);
            return false;
        }
    }

    /**
     * 移除会话
     */
    private void removeSession(WebSocketSession session) {
        Long userId = getUserIdFromSession(session);
        if (userId != null) {
            userSessions.remove(userId);
            log.info("移除用户 {} 的WebSocket会话", userId);
        }
    }

    /**
     * 根据会话对象移除会话
     */
    private void removeSessionBySession(WebSocketSession session) {
        userSessions.entrySet().removeIf(entry -> entry.getValue().equals(session));
    }

    /**
     * 获取当前连接数
     */
    public int getConnectionCount() {
        return userSessions.size();
    }

    /**
     * 检查用户是否在线
     */
    public boolean isUserOnline(Long userId) {
        WebSocketSession session = userSessions.get(userId);
        return session != null && session.isOpen();
    }
}
